# 恩格码密码机完成

class Enigma:
    # 初始化时，需要指定恩格尔密码机的三层偏移量
    def __init__(self, num1, num2, num3):
        # 一层转子：下面的列表作为26个字母的映射(一层),最后一个元素为偏转量
        self.ls1 = ['q', 'a', 'e', 'z', 'y', 'u', 'o', 's', 'c', 'v', 'p', 'l', 'i', 'h', 'w', 'b', 'n', 'k', 'r', 'g',
                    't', 'j', 'f', 'd', 'x', 'm', 0]
        # 二层转子：下面的列表作为26个字母的映射(一层),最后一个元素为偏转量
        self.ls2 = ['a', 'r', 'y', 'i', 'd', 'w', 'g', 'h', 'p', 'v', 'c', 'q', 'z', 's', 'x', 'n', 'b', 'm', 'l', 'k',
                    'e', 'o', 'f', 't', 'j', 'u', 0]
        # 三层转子：下面的列表作为26个字母的映射(一层),最后一个元素为偏转量
        self.ls3 = ['s', 'f', 'y', 'q', 'r', 'u', 'i', 'x', 'a', 'z', 'p', 'b', 'j', 'l', 'd', 'c', 'e', 'm', 'n', 'w',
                    'o', 'v', 't', 'g', 'h', 'k', 0]

        # 反射板 : 这是为了保证恩格玛机输入字符可逆的必要操作
        self.map4 = {'a': 'x', 'x': 'a',
                     'q': 'w', 'w': 'q',
                     's': 'f', 'f': 's',
                     'o': 'g', 'g': 'o',
                     'z': 'n', 'n': 'z',
                     'm': 't', 't': 'm',
                     'u': 'd', 'd': 'u',
                     'v': 'r', 'r': 'v',
                     'l': 'j', 'j': 'l',
                     'y': 'i', 'i': 'y',
                     'e': 'c', 'c': 'e',
                     'k': 'b', 'b': 'k',
                     'p': 'h', 'h': 'p'
                     }
        self.ls1[-1] = num1
        self.ls2[-1] = num2
        self.ls3[-1] = num3

    # 对偏移量的自增运算
    def up_num(self, num):  # num为偏移量
        num += 1
        if num >= 26:
            num = 0
        return num

    # 对输入的字符与偏移量的定位运算
    def get_index(self, inchar, num):  # inchar为输入字符，num为偏移量
        num += ord(inchar) - ord('a')
        if num > 25:
            num -= 26
        return num

    # 逆替换运算
    def get_index_change(self, inchar, ls):  # inchar为输入字符，ls为当前逆替换运算的转子
        num = ls.index(inchar) - ls[-1]
        if num < 0:
            num += 26
        c = chr(ord('a') + num)
        return c

    # 对一个字符串进行输入，分解为每个字符进行条件判断处理
    def instr(self, s):
        outt = []
        for i in s:
            if 0 <= ord(i) - ord('a') <= 25:  # 如果是小写字母a-z
                outt.append(self.inchar(i))  # 字母变换后存入列表
            elif 0 <= ord(i) - ord('A') <= 25:  # 如果是大写字母A-Z
                change = ord('a') - ord('A')
                ii = chr(ord(i) + change)  # 大写字母转小写
                outt.append(chr(ord(self.inchar(ii)) - change))  # 小写字母变回大写并存入列表
            else:
                outt.append(i)
        return ''.join(outt)  # !!!!!!!!!!!!!!!!!是否可以返回？

    # 核心功能，字母替换
    def inchar(self, c):
        # 第一层替换
        c = self.ls1[self.get_index(c, self.ls1[-1])]
        # 第二层替换
        c = self.ls2[self.get_index(c, self.ls2[-1])]
        # 第三层替换
        c = self.ls3[self.get_index(c, self.ls3[-1])]
        # 反射板反射
        c = self.map4[c]
        # 第三层逆替换（注意逻辑和普通替换不同）
        c = self.get_index_change(c, self.ls3)
        # 第二层逆替换（注意逻辑和普通替换不同）
        c = self.get_index_change(c, self.ls2)
        # 第一层逆替换（注意逻辑和普通替换不同）
        c = self.get_index_change(c, self.ls1)
        # 转子转动，偏移量自增
        self.ls1[-1] = self.up_num(self.ls1[-1])
        self.ls2[-1] = self.up_num(self.ls2[-1])
        self.ls3[-1] = self.up_num(self.ls3[-1])
        return c



a = Enigma(0, 1, 5) #这三个参数，分别为每层转子的偏移量，如需解码，需要设置偏移量和加密的时候一样即可。
print(a.instr("Siijx Vairt!"))
print(a.instr("Lzpsx Elljl!"))
print(a.instr("Lpcrt Ojnit!"))
